import json

import streamlit as st
import requests

# streamlit run E:\project\ai-chat\project\web311.py
# 设置页面标题
st.title("智能聊天AI")

# 初始化聊天历史
if "messages" not in st.session_state:
    st.session_state.messages = []

# 显示聊天历史
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

print("初始化聊天历史记录")
# 创建输入框和发送按钮
user_input = st.chat_input("请输入你的消息")
if user_input:
    if user_input:
        print("user_input:" + user_input)
        # 将用户输入添加到聊天历史记录
        st.session_state.messages.append({"role": "user", "content": user_input})
        with st.chat_message("user"):
            st.markdown(user_input)

        # 构建请求数据
        data = {
            "input": {
                "query": user_input
            }
        }

        # 发起流式请求
        url = "http://localhost:8000/topic/stream"
        response = requests.post(url, json=data, stream=True)
        if response.status_code == 200:
            answer = ""
            # 初始化一个占位符用于更新机器人消息
            with st.chat_message("assistant"):
                placeholder = st.empty()
            # 为机器人消息创建一个初始记录
            st.session_state.messages.append({"role": "assistant", "content": answer})
            line = ""
            for line in response.iter_lines():
                if line:
                    # 去除 SSE 数据的前缀 "data: "
                    line = line.lstrip(b'data: ').decode('utf-8')
                    if line.startswith("{"):
                        data = json.loads(line)
                        line = data.get('result')
                        if line:
                            answer += line
                            # 更新最新的机器人消息记录
                            st.session_state.messages[-1]["content"] = answer
                            # 更新页面显示
                            placeholder.markdown(f"<p style='text-align: left; background-color: #ffffe0; padding: 10px; border-radius: 5px;'>{answer}</p>",
                                                 unsafe_allow_html=True)

        else:
            st.error(f"请求失败，状态码: {response.status_code}")
